# client.py
import models
import copy
import torch
import torch.nn.functional
import torch.utils.data


class Client(object):

	def __init__(self, conf, model, train_dataset, id=-1):
		
		self.conf = conf
		
		self.local_model = model
		
		self.client_id = id
		
		self.train_dataset = train_dataset
		
		all_range = list(range(len(self.train_dataset)))
		data_len = int(len(self.train_dataset) / self.conf['no_models'])
		train_indices = all_range[id * data_len: (id + 1) * data_len]

		self.train_loader = torch.utils.data.DataLoader(self.train_dataset, batch_size=conf["batch_size"], sampler=torch.utils.data.sampler.SubsetRandomSampler(train_indices))
		# 无放回地按照给定的索引列表 采样 样本元素

	def local_train(self, model):

		for name, param in model.state_dict().items():  # 遍历模型参数
			self.local_model.state_dict()[name].copy_(param.clone())
	
		# print(id(model))
		# 使用随机梯度下降算法作为优化算法
		optimizer = torch.optim.SGD(self.local_model.parameters(), lr=self.conf['lr'], momentum=self.conf['momentum'])
		# print(id(self.local_model))
		self.local_model.train()
		for e in range(self.conf["local_epochs"]):
			
			for batch_id, batch in enumerate(self.train_loader):
				data, target = batch
				
				if torch.cuda.is_available():  # 将训练数据放到GPU环境
					data = data.cuda()
					target = target.cuda()
			
				optimizer.zero_grad()  # 梯度清零
				output = self.local_model(data)  # 前向传播算法
				loss = torch.nn.functional.cross_entropy(output, target)  # 使用交叉熵计算损失值
				loss.backward()  # 反向传播计算得到了梯度
			
				optimizer.step()  # 利用反向传播得到的梯度，利用优化算法更新网络参数（权重）
			print("Local Epoch %d done." % (e+1))
		diff = dict()  # 生成一个空的字典
		for name, data in self.local_model.state_dict().items():
			diff[name] = (data - model.state_dict()[name])
			# print(diff[name])
			
		return diff  # 返回网络参数的变化
		